import numpy as np

states = [[1, 0, 1], [0, 1, 0], [1, 1, 0]]
states = [states]
states = np.array(states)
distributed = np.count_nonzero(states, axis=2) > 1 # length = #jobs, each entry is the used node number of each job.
print(distributed)
distributed = np.expand_dims(distributed, axis=-1) 
mask = states * distributed
print(states, distributed, mask)
mask = mask > 0
mask = mask.cumsum(axis=1) > 1
print(mask)
states[mask] = 0
print(states)
